package com.diaohw.platform.common.util;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;

/***
 * 树形结构
 */
public class TreeUtil {

    /**
     * 集合转tree
     *
     * @param list     集合
     * @param id       ID function
     * @param parent   parent function
     * @param children children function
     */
    public static <T, R> List<T> listToTree(List<T> list, Function<T, R> id, Function<T, R> parent, Function<T, List<T>> children) {

        Map<R, T> map = new HashMap<>();
        List<T> tree = new ArrayList<>();

        list.forEach(l -> {
            map.put(id.apply(l), l);
        });
        list.forEach(l -> {
            if (map.containsKey(parent.apply(l))) {
                T data = map.get(parent.apply(l));
                children.apply(data).add(l);
            } else {
                tree.add(l);
            }
        });

        return tree;

    }

    public static <T> List<T> treeToList(List<T> trees, Function<T, List<T>> children) {
        List<T> list = new ArrayList<>();
        for (T tree : trees) {
            list.add(tree);
            List<T> apply = children.apply(tree);
            if (apply == null || apply.isEmpty()) {
                continue;
            }
            List<T> c = treeToList(apply, children);
            apply.clear();
            list.addAll(c);
        }
        return list;
    }

}
